import time
from fastapi import APIRouter, WebSocket
from src.celery.ai_image import start_ai_img
from src.controller.ai.schemas import textSearchVideoSchemas
from src.config import *
from src.utils import *

router = APIRouter()


# http://127.0.0.1:8000/api/v1/ai/text_search_video
@router.post("/text_search_video")
async def textSearchVideo(params: textSearchVideoSchemas):
    task = start_ai_img.delay(params.keyWord, params.imgList)
    return { 
        'task_id': task.id,
        'message': 'ai任务已添加后台'
    }

# @router.websocket("/ws")
# async def task_status(websocket: WebSocket):
#     await websocket.accept()
#     while True:
#         data = await websocket.receive_json()
#         while True:
#             res = getTaskStaus(data['task_id'])
#             print("2 循环查找任务 res：", res)
#             if res['status']:
#                 await websocket.send_json(res)
#                 break;
#             time.sleep(3)
        

# http://127.0.0.1:8000/api/v1/ai/task_status?task_id=c4b57ff1-0ca4-448e-ad95-e3d549841842
@router.get("/task_status")
def taskStatus(task_id: str):
    return getTaskStaus(task_id)
    

def getTaskStaus(id: str):
    async_result = start_ai_img.AsyncResult(id)
    if async_result.successful():
        return async_result.get() | { 'status': True }
    else:
        return {
            'status': False,
            'message': tool.getState(async_result.status)
        }